package org.dreamwork.sso.saml2.servlet;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.dreamwork.misc.Base64;
import org.dreamwork.sso.saml2.SamlKeys;
import org.dreamwork.sso.saml2.protocol.*;
import org.dreamwork.sso.saml2.spec.IKeyFetcher;
import org.dreamwork.secure.AlgorithmMapping;
import org.dreamwork.sso.saml2.util.ConfigKey;
import org.dreamwork.util.StringUtil;
import org.joda.time.DateTime;
import org.opensaml.xml.*;

import javax.servlet.RequestDispatcher;
import javax.servlet.ServletContext;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.IOException;
import java.net.URLEncoder;
import java.security.*;
import java.security.cert.X509Certificate;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.UUID;

/**
 * Created with IntelliJ IDEA.
 * User: seth.yang
 * Date: 13-8-12
 * Time: 下午4:35
 */
public class SamlLoginHandler extends HttpServlet {
    private static final String SAML_REQUEST_TEMP_KEY = "saml.request.temp.key";
    private Properties props;
    private IKeyFetcher fetcher;
    private static final String ENC = "utf-8";

    private static final String PN_SAML_REQUEST = "SAMLRequest";
    private static final String PN_SIGNATURE_ALGORITHM = "SigAlg";
    private static final String PN_SIGNATURE = "Signature";

    @Override
    public void init () throws ServletException {
        ServletContext context = getServletContext ();
        props = (Properties) context.getAttribute (SamlKeys.CONFIG_CONTEXT_KEY);
        fetcher = (IKeyFetcher) context.getAttribute (SamlKeys.KEY_FETCHER_CONTEXT_KEY);
    }

    // url + "?SamlLoginRequest=" + samlRequest + "&SigAlg=" + sigAlg + "&Signature=" + signature;
    @SuppressWarnings ("unchecked")
    @Override
    protected void service (HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
        HttpSession session = request.getSession ();
        Object entity = session.getAttribute (SamlKeys.SAML_ENTITY_SESSION_KEY);

        if (entity == null) {
            // 第一次访问
            processFirstRequest (request, response);
        } else/* if (status == SamlStatus.SUCCESSFUL) */{
            try {
                // 登录成功，有两种可能 1.当前 SP 引导登录成功, 2. 其他 SP 曾经引导登录成功
                String s_saml_request, algorithm, s_signature;
                Map<String, String> temp = (Map<String, String>) session.getAttribute (SAML_REQUEST_TEMP_KEY);
                if (temp != null) { // 从登录页面返回
                    s_saml_request = temp.get (PN_SAML_REQUEST);
                    algorithm = temp.get (PN_SIGNATURE_ALGORITHM);
                    s_signature = temp.get (PN_SIGNATURE);
                } else { // 其他 SP 请求登录
                    s_saml_request = request.getParameter (PN_SAML_REQUEST);
                    algorithm = request.getParameter (PN_SIGNATURE_ALGORITHM);
                    s_signature = request.getParameter (PN_SIGNATURE);
                }

                // 构建 SAML Login Request
                SamlLoginRequest samlLoginRequest = SamlLoginRequest.parse (s_saml_request);

                // 针对当前 SP，准备校验请求数据
                String issuer = samlLoginRequest.getIssuer ();
                ServletContext context = getServletContext ();
                PublicKey publicKey = fetcher.getPublicKey (context, issuer);
                if (publicKey == null) {
                    response.sendError (
                            HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
                            "There's no public key association to service provider [" + issuer + "]"
                    );

                    return;
                }

                // 构建空 SAML Login Response
                SamlLoginResponse samlLoginResponse = createEmptyResponse ();

                if (!samlLoginRequest.validate (publicKey, s_signature, algorithm)) {
                    // 校验失败，设置返回错误信息
                    samlLoginResponse.setStatus (SamlStatus.INVALID_SIGNATURE);
                } else {
                    SamlStatus status = (SamlStatus) session.getAttribute (SamlKeys.SAML_LOGIN_STATUS);
                    samlLoginResponse.setStatus (status);
                    String uri = getIdpURI (request);
                    // 创建 SAML 断言, 包裹认证凭证
                    SamlAssertion assertion = buildAssertion (uri, issuer, entity);
                    samlLoginResponse.setAssertion (assertion);

                    X509Certificate certificate = fetcher.getCertificate (context, issuer);
                    PrivateKey privateKey = fetcher.getPrivateKey (context, issuer);
                    // 加密断言
                    samlLoginResponse.encrypt (publicKey);
                    // 签署响应
                    samlLoginResponse.sign (certificate, privateKey);
                }

                perform (request, response, samlLoginRequest, samlLoginResponse);
            } catch (Exception ex) {
                Log logger = LogFactory.getLog (SamlLoginHandler.class);
                logger.warn (ex.getMessage (), ex);
                session.removeAttribute (SamlKeys.SAML_LOGIN_STATUS);
                throw new ServletException (ex);
            } finally {
                session.removeAttribute (SAML_REQUEST_TEMP_KEY);
            }
        }
    }

    private SamlAssertion buildAssertion (String idpURI, String issuer, Object entity) {
        XMLObjectBuilderFactory factory = Configuration.getBuilderFactory ();
        SamlSubject subject = new SamlSubject (factory);
        subject.addSubjectConfirmation ("urn:oasis:names:tc:SAML:2.0:cm:bearer");
        subject.setNameId ('_' + UUID.randomUUID ().toString ());

        SamlAssertion assertion = new SamlAssertion (factory);
        assertion.setNotOnOrAfter (new DateTime ());
        assertion.setId ('_' + UUID.randomUUID ().toString ());
        assertion.setAuthnContext (idpURI, '_' + UUID.randomUUID ().toString (), 10);
        assertion.setIssueInstant (new DateTime ());
        assertion.setSubject (subject);
        assertion.setIssuer (issuer);
        assertion.setAttributes (entity);

        return assertion;
    }

    private String getIdpURI (HttpServletRequest request) {
        String protocol = request.getProtocol ();
        String path = request.getServerName ();
        int port = request.getServerPort ();

        StringBuilder builder = new StringBuilder ();
        builder.append (protocol).append ("://").append (path);
        if ("https".equalsIgnoreCase (protocol) && port == 443) {

        } else if ("http".equalsIgnoreCase (protocol) && port == 80) {

        } else {
            builder.append (":").append (port);
        }
        return builder.toString ();
    }

    private void perform (
            HttpServletRequest request, HttpServletResponse response,
            SamlLoginRequest samlLoginRequest, SamlLoginResponse samlLoginResponse
    ) throws IOException, NoSuchAlgorithmException, InvalidKeyException, SignatureException, ServletException {
        String action = samlLoginRequest.getAssertionConsumerServiceURL ();
        String uri = AlgorithmMapping.Signature.RSA_SHA1.uri;

        String xml = samlLoginResponse.toXMLString ();
        String content = "SAMLResponse=" + URLEncoder.encode (xml, ENC) + "&SigAlg=" + URLEncoder.encode (uri, ENC);
        byte[] rawData = content.getBytes ();
        String jceName = AlgorithmMapping.translateSignatureAlgorithm (uri);
        Signature signature = Signature.getInstance (jceName);
        PrivateKey key = fetcher.getPrivateKey (getServletContext (), samlLoginRequest.getIssuer ());
        signature.initSign (key);
        signature.update (rawData);
        byte[] buff = signature.sign ();
        byte[] keyInfo = samlLoginResponse.getSignatureKeyInfo ();

        request.setAttribute ("action", action);
        request.setAttribute ("sigAlg", uri);
        request.setAttribute ("signature", new String (Base64.encode (buff), ENC));
        request.setAttribute ("samlResponse", new String (Base64.encode (xml.getBytes (ENC)), ENC));
        request.setAttribute ("keyInfo", new String (Base64.encode (keyInfo), ENC));

        String target = props.getProperty (ConfigKey.SAML2_PERFORM_PAGE);
        RequestDispatcher dispatcher = request.getRequestDispatcher (target);
        dispatcher.forward (request, response);
    }

    private void processFirstRequest (HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
        String s_saml_request = request.getParameter (PN_SAML_REQUEST);
        String algorithm = request.getParameter (PN_SIGNATURE_ALGORITHM);
        String s_signature = request.getParameter (PN_SIGNATURE);

        if (StringUtil.isEmpty (s_saml_request)) {
            response.sendError (HttpServletResponse.SC_BAD_REQUEST, "SamlLoginRequest is null or empty");
            return;
        }
        if (StringUtil.isEmpty (algorithm)) {
            response.sendError (HttpServletResponse.SC_BAD_REQUEST, "Signature Algorithm is null or empty");
            return;
        }
        if (StringUtil.isEmpty (s_signature)) {
            response.sendError (HttpServletResponse.SC_BAD_REQUEST, "Signature is null or empty");
            return;
        }

        Map<String, String> temp = new HashMap<String, String> ();
        temp.put (PN_SAML_REQUEST, s_saml_request);
        temp.put (PN_SIGNATURE_ALGORITHM, algorithm);
        temp.put (PN_SIGNATURE, s_signature.replace (' ', '+'));
        request.getSession ().setAttribute (SAML_REQUEST_TEMP_KEY, temp);

        String target = props.getProperty (ConfigKey.SAML2_LOGIN_PAGE);
        RequestDispatcher dispatcher = request.getRequestDispatcher (target);
        dispatcher.forward (request, response);
    }

    private SamlLoginResponse createEmptyResponse () {
        XMLObjectBuilderFactory factory = Configuration.getBuilderFactory ();
        SamlLoginResponse samlLoginResponse = new SamlLoginResponse (factory);
        samlLoginResponse.setIssuer (props.getProperty (ConfigKey.SAML2_ISSUER));

        return samlLoginResponse;
    }
}